from typing import Annotated, Sequence, List
from PIL import Image
from marker.processors import BaseProcessor
from marker.schema import BlockTypes
from marker.schema.blocks import Block, BlockId
from marker.schema.document import Document
from marker.logger import get_logger

logger = get_logger()

class BaseOptProcessor(BaseProcessor):
    """
    Block识别优化
    """
    max_concurrency: Annotated[
        int,
        "The maximum number of concurrent requests to make to the Gemini model.",
    ] = 5
    image_expansion_ratio: Annotated[
        float,
        "The ratio to expand the image by when cropping.",
    ] = 0.01
    disable_tqdm: Annotated[
        bool,
        "Whether to disable the tqdm progress bar.",
    ] = False
    block_types = None

    def __init__(self, config=None):
        super().__init__(config)

    def extract_image(
        self,
        document: Document,
        image_block: Block,
        remove_blocks: Sequence[BlockTypes] | None = None,
    ) -> Image.Image:
        return image_block.get_image(
            document,
            highres=True,
            expansion=(self.image_expansion_ratio, self.image_expansion_ratio),
            remove_blocks=remove_blocks,
        )

    @staticmethod
    def handle_rewrites(blocks: list, document: Document):
        for block_data in blocks:
            try:
                block_id = block_data["id"].strip().lstrip("/")
                _, page_id, block_type, block_id = block_id.split("/")
                block_id = BlockId(
                    page_id=page_id,
                    block_id=block_id,
                    block_type=getattr(BlockTypes, block_type),
                )
                block = document.get_block(block_id)
                if not block:
                    logger.debug(f"Block {block_id} not found in document")
                    continue

                if hasattr(block, "html"):
                    block.html = block_data["html"]
            except Exception as e:
                logger.debug(f"Error parsing block ID {block_data['id']}: {e}")
                continue

    def __call__(self, document: Document):
        try:
            doc_blocks = []
            for page in document.pages:
                page_blocks = page.contained_blocks(document, self.block_types)
                page_blocks = list(filter(lambda item: item.block_type in self.block_types, page_blocks))
                doc_blocks.extend(page_blocks)
            if len(doc_blocks) == 0:
                return
            block_images = []
            for block in doc_blocks:
                image = self.extract_image(document, block)
                block_images.append(image)
            doc_blocks = self.process_rewriting(doc_blocks, block_images)
            self.handle_rewrites(doc_blocks, document)
        except Exception as e:
            logger.warning(f"Error rewriting blocks in {self.__class__.__name__}: {e}")

    def process_rewriting(self, doc_blocks: List[Block], block_images: List[Image.Image]):
        raise NotImplementedError()
